from Meta_RL.src.utils.serializable import Serializable
from Meta_RL.src.dynamics.core.utils import create_rnn
import tensorflow as tf
import numpy as np


class RNNDynamicsModel(Serializable):
    def __init__(self,
                name,
                env,
                hidden_sizes=(512,),
                cell_type = 'lstm',
                optimizer=tf.train.AdamOptimizer,
                valid_split_ratio=0.2,
                learning_rate=.001
                ):
        Serializable.quick_init(self, locals())

        self.valid_split_ratio = valid_split_ratio

        self.learning_rate = learning_rate

        self.obs_space_dims = obs_space_dims = env.observation_space.shape[0]
        self.act_space_dims = act_space_dims = env.action_space.shape[0]

        with tf.variable_scope(name):
            self.obs_ph = tf.placeholder(tf.float32, shape=(None, None, obs_space_dims))
            self.act_ph = tf.placeholder(tf.float32, shape=(None, None, act_space_dims))
            self.delta_ph = tf.placeholder(tf.float32, shape=(None, None, obs_space_dims))# 预测的delta obs

            self.nn_input = tf.concat([self.obs_ph, self.act_ph], axis=2)

            with tf.variable_scope('rnn_model'):
                rnn = 


